#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基金净值查询压力测试脚本

功能:
1. 统计fund_nav_detail表中的基金总数和记录总数
2. 随机选择基金进行净值查询
3. 测量查询时长并进行性能分析
4. 基于测试结果提供优化建议

Author: AI Assistant
Date: 2025-01-27
"""

import os
import sys
import time
import random
import logging
import statistics
from typing import Dict, List, Any
from datetime import datetime, timedelta

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
backend_dir = os.path.dirname(current_dir)
sys.path.insert(0, backend_dir)

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'wealth_backend.settings')

try:
    import django  # type: ignore
    django.setup()
except ImportError as e:
    print(f"Django导入错误: {e}")
    print("请确保已安装Django并且在正确的虚拟环境中运行")
    sys.exit(1)

# 导入Django数据库连接
from django.db import connection  # type: ignore
from data_service.models import FundNavDetail  # type: ignore

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class FundNavPerformanceTester:
    """
    基金净值查询性能测试器
    """
    
    def __init__(self):
        """初始化测试器"""
        self.test_results = {
            'total_funds': 0,
            'total_records': 0,
            'query_times': [],
            'test_samples': [],
            'errors': []
        }
        logger.info("基金净值性能测试器初始化完成")
    
    def test_database_connection(self) -> bool:
        """
        测试数据库连接
        
        Returns:
            bool: 连接是否成功
        """
        try:
            with connection.cursor() as cursor:
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                if result:
                    logger.info("数据库连接成功")
                    return True
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            self.test_results['errors'].append(f"数据库连接失败: {e}")
        return False
    
    def get_table_statistics(self) -> Dict[str, Any]:
        """
        获取fund_nav_detail表的基本统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        stats = {
            'total_funds': 0,
            'total_records': 0,
            'date_range': {},
            'sample_funds': []
        }
        
        try:
            with connection.cursor() as cursor:
                # 统计基金总数
                cursor.execute("""
                    SELECT COUNT(DISTINCT fund_code) as fund_count
                    FROM fund_NetAssetValue
                """)
                result = cursor.fetchone()
                stats['total_funds'] = result[0] if result else 0
                
                # 统计记录总数
                cursor.execute("""
                    SELECT COUNT(*) as record_count
                    FROM fund_NetAssetValue
                """)
                result = cursor.fetchone()
                stats['total_records'] = result[0] if result else 0
                
                # 获取日期范围
                cursor.execute("""
                    SELECT 
                        MIN(nav_date) as earliest_date,
                        MAX(nav_date) as latest_date
                    FROM fund_NetAssetValue
                """)
                result = cursor.fetchone()
                if result:
                    stats['date_range'] = {
                        'earliest': result[0],
                        'latest': result[1]
                    }
                
                # 获取样本基金列表（随机选择50个基金）
                cursor.execute("""
                    SELECT DISTINCT fund_code
                    FROM fund_NetAssetValue
                    ORDER BY RAND()
                    LIMIT 50
                """)
                results = cursor.fetchall()
                stats['sample_funds'] = [row[0] for row in results]
            
            self.test_results['total_funds'] = stats['total_funds']
            self.test_results['total_records'] = stats['total_records']
            
            logger.info(f"表统计完成: {stats['total_funds']}个基金, {stats['total_records']}条记录")
            
        except Exception as e:
            logger.error(f"获取表统计信息失败: {e}")
            self.test_results['errors'].append(f"获取表统计信息失败: {e}")
        
        return stats
    
    def test_single_fund_query(self, fund_code: str, query_type: str = 'recent') -> Dict[str, Any]:
        """
        测试单个基金的查询性能
        
        Args:
            fund_code: 基金代码
            query_type: 查询类型 ('recent', 'all', 'range')
            
        Returns:
            Dict[str, Any]: 查询结果和性能数据
        """
        result = {
            'fund_code': fund_code,
            'query_type': query_type,
            'query_time': 0,
            'record_count': 0,
            'success': False,
            'error': None
        }
        
        try:
            with connection.cursor() as cursor:
                # 根据查询类型构建不同的SQL
                if query_type == 'recent':
                    # 查询最近30天的数据
                    sql = """
                        SELECT fund_code, nav_date, unit_net_value, accumulated_net_value, daily_growth_rate
                        FROM fund_NetAssetValue
                        WHERE fund_code = %s
                        AND nav_date >= DATE_SUB(CURDATE(), INTERVAL 30 DAY)
                        ORDER BY nav_date DESC
                    """
                    params = (fund_code,)
                elif query_type == 'all':
                    # 查询所有历史数据
                    sql = """
                        SELECT fund_code, nav_date, unit_net_value, accumulated_net_value, daily_growth_rate
                        FROM fund_NetAssetValue
                        WHERE fund_code = %s
                        ORDER BY nav_date DESC
                    """
                    params = (fund_code,)
                elif query_type == 'range':
                    # 查询指定日期范围的数据（最近1年）
                    sql = """
                        SELECT fund_code, nav_date, unit_net_value, accumulated_net_value, daily_growth_rate
                        FROM fund_NetAssetValue
                        WHERE fund_code = %s
                        AND nav_date BETWEEN DATE_SUB(CURDATE(), INTERVAL 1 YEAR) AND CURDATE()
                        ORDER BY nav_date DESC
                    """
                    params = (fund_code,)
                else:
                    raise ValueError(f"不支持的查询类型: {query_type}")
                
                # 执行查询并测量时间
                start_time = time.time()
                cursor.execute(sql, params)
                results = cursor.fetchall()
                end_time = time.time()
                
                result['query_time'] = end_time - start_time
                result['record_count'] = len(results)
                result['success'] = True
            
        except Exception as e:
            result['error'] = str(e)
            logger.error(f"查询基金 {fund_code} 失败: {e}")
            self.test_results['errors'].append(f"查询基金 {fund_code} 失败: {e}")
        
        return result
    
    def run_performance_test(self, sample_size: int = 20, test_rounds: int = 3) -> Dict[str, Any]:
        """
        运行性能测试
        
        Args:
            sample_size: 测试样本数量
            test_rounds: 测试轮数
            
        Returns:
            Dict[str, Any]: 测试结果
        """
        logger.info(f"开始性能测试: {sample_size}个样本, {test_rounds}轮测试")
        
        # 获取表统计信息
        stats = self.get_table_statistics()
        
        if not stats['sample_funds']:
            logger.error("没有可用的测试基金")
            return self.test_results
        
        # 随机选择测试基金
        test_funds = random.sample(
            stats['sample_funds'], 
            min(sample_size, len(stats['sample_funds']))
        )
        
        query_types = ['recent', 'range', 'all']
        
        for round_num in range(test_rounds):
            logger.info(f"执行第 {round_num + 1} 轮测试")
            
            for fund_code in test_funds:
                for query_type in query_types:
                    # 添加随机延时避免过于频繁的查询
                    time.sleep(random.uniform(0.1, 0.3))
                    
                    result = self.test_single_fund_query(fund_code, query_type)
                    
                    if result['success']:
                        self.test_results['query_times'].append(result['query_time'])
                        self.test_results['test_samples'].append(result)
                        
                        logger.info(
                            f"基金 {fund_code} ({query_type}): "
                            f"{result['query_time']:.4f}s, {result['record_count']}条记录"
                        )
        
        return self.test_results
    
    def analyze_results(self) -> Dict[str, Any]:
        """
        分析测试结果
        
        Returns:
            Dict[str, Any]: 分析结果
        """
        if not self.test_results['query_times']:
            return {'error': '没有有效的测试数据'}
        
        query_times = self.test_results['query_times']
        
        analysis = {
            'performance_metrics': {
                'total_queries': len(query_times),
                'avg_query_time': statistics.mean(query_times),
                'median_query_time': statistics.median(query_times),
                'min_query_time': min(query_times),
                'max_query_time': max(query_times),
                'std_deviation': statistics.stdev(query_times) if len(query_times) > 1 else 0
            },
            'query_type_analysis': {},
            'performance_rating': '',
            'bottlenecks': [],
            'optimization_suggestions': []
        }
        
        # 按查询类型分析
        for query_type in ['recent', 'range', 'all']:
            type_samples = [s for s in self.test_results['test_samples'] if s['query_type'] == query_type]
            if type_samples:
                type_times = [s['query_time'] for s in type_samples]
                analysis['query_type_analysis'][query_type] = {
                    'count': len(type_times),
                    'avg_time': statistics.mean(type_times),
                    'avg_records': statistics.mean([s['record_count'] for s in type_samples])
                }
        
        # 性能评级
        avg_time = analysis['performance_metrics']['avg_query_time']
        if avg_time < 0.1:
            analysis['performance_rating'] = '优秀'
        elif avg_time < 0.5:
            analysis['performance_rating'] = '良好'
        elif avg_time < 1.0:
            analysis['performance_rating'] = '一般'
        else:
            analysis['performance_rating'] = '需要优化'
        
        # 识别性能瓶颈
        if avg_time > 0.5:
            analysis['bottlenecks'].append('平均查询时间过长')
        
        if analysis['performance_metrics']['max_query_time'] > 2.0:
            analysis['bottlenecks'].append('存在极慢查询')
        
        if analysis['performance_metrics']['std_deviation'] > 0.5:
            analysis['bottlenecks'].append('查询时间波动较大')
        
        return analysis
    
    def generate_optimization_suggestions(self, analysis: Dict[str, Any]) -> List[str]:
        """
        生成优化建议
        
        Args:
            analysis: 分析结果
            
        Returns:
            List[str]: 优化建议列表
        """
        suggestions = []
        
        avg_time = analysis['performance_metrics']['avg_query_time']
        total_records = self.test_results['total_records']
        
        # 基于平均查询时间的建议
        if avg_time > 1.0:
            suggestions.extend([
                "1. 数据库索引优化:",
                "   - 确保fund_code字段有索引",
                "   - 在(fund_code, nav_date)上创建复合索引",
                "   - 考虑在nav_date字段上创建单独索引"
            ])
        
        if avg_time > 0.5:
            suggestions.extend([
                "2. 查询优化:",
                "   - 使用LIMIT限制返回结果数量",
                "   - 避免SELECT *，只查询需要的字段",
                "   - 考虑使用分页查询"
            ])
        
        # 基于数据量的建议
        if total_records > 1000000:  # 超过100万条记录
            suggestions.extend([
                "3. 数据分区优化:",
                "   - 考虑按日期对表进行分区",
                "   - 实施数据归档策略，移除过期数据",
                "   - 使用读写分离架构"
            ])
        
        # 基于查询类型的建议
        if 'query_type_analysis' in analysis:
            all_query = analysis['query_type_analysis'].get('all', {})
            if all_query and all_query.get('avg_time', 0) > 2.0:
                suggestions.extend([
                    "4. 全量查询优化:",
                    "   - 避免一次性查询所有历史数据",
                    "   - 实施懒加载或分批加载",
                    "   - 考虑使用缓存机制"
                ])
        
        # 通用优化建议
        suggestions.extend([
            "5. 系统级优化:",
            "   - 增加数据库连接池大小",
            "   - 优化MySQL配置参数",
            "   - 考虑使用Redis缓存热点数据",
            "   - 监控慢查询日志"
        ])
        
        # 应用层优化
        suggestions.extend([
            "6. 应用层优化:",
            "   - 实施查询结果缓存",
            "   - 使用异步查询处理",
            "   - 批量查询多个基金数据",
            "   - 实施数据预加载策略"
        ])
        
        return suggestions
    
    def print_detailed_report(self):
        """
        打印详细的测试报告
        """
        print("\n" + "="*80)
        print("基金净值查询性能测试报告")
        print("="*80)
        print(f"测试时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        
        # 基本统计信息
        print("\n📊 数据库统计信息:")
        print("-"*50)
        print(f"基金总数: {self.test_results['total_funds']:,}")
        print(f"记录总数: {self.test_results['total_records']:,}")
        
        # 分析结果
        analysis = self.analyze_results()
        
        if 'error' in analysis:
            print(f"\n❌ 分析失败: {analysis['error']}")
            return
        
        # 性能指标
        metrics = analysis['performance_metrics']
        print("\n⚡ 性能指标:")
        print("-"*50)
        print(f"总查询次数: {metrics['total_queries']}")
        print(f"平均查询时间: {metrics['avg_query_time']:.4f}秒")
        print(f"中位数查询时间: {metrics['median_query_time']:.4f}秒")
        print(f"最快查询时间: {metrics['min_query_time']:.4f}秒")
        print(f"最慢查询时间: {metrics['max_query_time']:.4f}秒")
        print(f"标准差: {metrics['std_deviation']:.4f}秒")
        print(f"性能评级: {analysis['performance_rating']}")
        
        # 查询类型分析
        if analysis['query_type_analysis']:
            print("\n📈 查询类型分析:")
            print("-"*50)
            for query_type, data in analysis['query_type_analysis'].items():
                type_name = {
                    'recent': '最近30天查询',
                    'range': '最近1年查询', 
                    'all': '全量历史查询'
                }.get(query_type, query_type)
                print(f"{type_name}:")
                print(f"  查询次数: {data['count']}")
                print(f"  平均时间: {data['avg_time']:.4f}秒")
                print(f"  平均记录数: {data['avg_records']:.0f}条")
        
        # 性能瓶颈
        if analysis['bottlenecks']:
            print("\n⚠️  识别的性能瓶颈:")
            print("-"*50)
            for bottleneck in analysis['bottlenecks']:
                print(f"• {bottleneck}")
        
        # 优化建议
        suggestions = self.generate_optimization_suggestions(analysis)
        if suggestions:
            print("\n💡 优化建议:")
            print("-"*50)
            for suggestion in suggestions:
                print(suggestion)
        
        # 错误信息
        if self.test_results['errors']:
            print("\n❌ 测试过程中的错误:")
            print("-"*50)
            for error in self.test_results['errors']:
                print(f"• {error}")
        
        print("\n" + "="*80)
        print("测试报告结束")
        print("="*80)

def main():
    """
    主函数
    """
    print("基金净值查询性能测试开始...")
    
    # 创建测试器
    tester = FundNavPerformanceTester()
    
    try:
        # 测试数据库连接
        if not tester.test_database_connection():
            print("数据库连接失败，测试终止")
            return
        
        # 运行性能测试
        # 可以调整参数：sample_size=测试基金数量, test_rounds=测试轮数
        results = tester.run_performance_test(sample_size=15, test_rounds=2)
        
        # 打印详细报告
        tester.print_detailed_report()
        
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    except Exception as e:
        logger.error(f"测试过程中发生未预期的错误: {e}")
        print(f"测试失败: {e}")
    finally:
        # 数据库连接会自动关闭
        print("\n测试完成")

if __name__ == '__main__':
    main()